{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "mHF9VCProKJN"
   },
   "source": [
    "# AI Explanations: Deploying an image model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "hZzRVxNtH-zG"
   },
   "source": [
    "## Overview\n",
    "\n",
    "This tutorial shows how to train a Keras classification model on image data and deploy it to the AI Platform Explanations service to get feature attributions on your deployed model.\n",
    "\n",
    "If you've already got a trained model and want to deploy it to AI Explanations, skip to the **Export the model as a TF 2 SavedModel** section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "iN69d4D9Flrh"
   },
   "source": [
    "### Dataset\n",
    "\n",
    "The dataset used for this tutorial is the [flowers dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "Su2qu-4CW-YH"
   },
   "source": [
    "### Objective\n",
    "\n",
    "The goal of this tutorial is to train a model on a simple image dataset (flower classification) to understand how you can use AI Explanations with image models. For image models, AI Explanations returns an image with the pixels highlighted that signaled your model's prediction the most.\n",
    "\n",
    "This tutorial focuses more on deploying the model to AI Platform with Explanations than on the design of the model itself. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "rgLXkyHEvTVD"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L4bqjh0QAMB8"
   },
   "outputs": [],
   "source": [
    "PROJECT_ID = \"[your-project-id]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-8HMrmgNAMB-"
   },
   "outputs": [],
   "source": [
    "if PROJECT_ID == \"\" or PROJECT_ID is None or PROJECT_ID == \"[your-project-id]\":\n",
    "    # Get your GCP project id from gcloud\n",
    "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
    "    PROJECT_ID = shell_output[0]\n",
    "    print(\"Project ID:\", PROJECT_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nXTGJUOfKHUe"
   },
   "outputs": [],
   "source": [
    "!gcloud config set project $PROJECT_ID"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fiFJqXOGAMCB"
   },
   "source": [
    "#### Region\n",
    "\n",
    "You can also change the `REGION` variable, which is used for operations\n",
    "throughout the rest of this notebook. Make sure to [choose a region where Cloud\n",
    "AI Platform services are\n",
    "available](https://cloud.google.com/ml-engine/docs/tensorflow/regions). You can\n",
    "not use a Multi-Regional Storage bucket for training with AI Platform."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aJboCICuAMCC"
   },
   "outputs": [],
   "source": [
    "REGION = 'us-central1'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "acproI4vAMCE"
   },
   "source": [
    "#### Timestamp\n",
    "\n",
    "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append onto the name of resources which will be created in this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "w5Khg6W9AMCE"
   },
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "tT061irlJwkg"
   },
   "source": [
    "### Create a Cloud Storage bucket\n",
    "\n",
    "When you submit a training job using the Cloud SDK, you upload a Python package\n",
    "containing your training code to a Cloud Storage bucket. AI Platform runs\n",
    "the code from this package. In this tutorial, AI Platform also saves the\n",
    "trained model that results from your job in the same bucket. You can then\n",
    "create an AI Platform model version based on this output in order to serve\n",
    "online predictions.\n",
    "\n",
    "Set the name of your Cloud Storage bucket below. It must be unique across all\n",
    "Cloud Storage buckets. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jrLa8bYQAMCJ"
   },
   "outputs": [],
   "source": [
    "BUCKET_NAME = \"[your-bucket-name]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "bTxmbDg1I0x1"
   },
   "outputs": [],
   "source": [
    "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"[your-bucket-name]\":\n",
    "    BUCKET_NAME = PROJECT_ID + \"_xai_flowers_\" + TIMESTAMP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "fsmCk2dwJnLZ"
   },
   "source": [
    "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "160PRO3aJqLD"
   },
   "outputs": [],
   "source": [
    "! gcloud storage buckets create --location=$REGION gs://$BUCKET_NAME"   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PyxoF-iqqD1t"
   },
   "source": [
    "### Import libraries\n",
    "\n",
    "Import the libraries for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install explainable-ai-sdk\n",
    "!pip install explainable-ai-sdk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MEDlLSWK15UL"
   },
   "outputs": [],
   "source": [
    "import io\n",
    "import os\n",
    "import random\n",
    "import numpy as np\n",
    "import PIL\n",
    "import tensorflow as tf\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "from base64 import b64encode\n",
    "\n",
    "# should be >= 2.1\n",
    "print(\"Tensorflow version \" + tf.__version__)\n",
    "if tf.__version__ < \"2.1\":\n",
    "    raise Exception(\"TF 2.1 or greater is required\")\n",
    "\n",
    "AUTO = tf.data.experimental.AUTOTUNE\n",
    "print(\"AUTO\", AUTO)\n",
    "\n",
    "import explainable_ai_sdk"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "aRVMEU2Qshm4"
   },
   "source": [
    "## Download and preprocess the data\n",
    "\n",
    "This section shows how to download the flower images, use the `tf.data` API to create a data input pipeline, and split the data into training and validation sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "v7HLNsvekxvz"
   },
   "outputs": [],
   "source": [
    "GCS_PATTERN = 'gs://flowers-public/tfrecords-jpeg-192x192-2/*.tfrec'\n",
    "IMAGE_SIZE = [192, 192]\n",
    "\n",
    "BATCH_SIZE = 32\n",
    "\n",
    "VALIDATION_SPLIT = 0.19\n",
    "CLASSES = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']  # do not change, maps to the labels in the data (folder names)\n",
    "\n",
    "# Split data files between training and validation\n",
    "filenames = tf.io.gfile.glob(GCS_PATTERN)\n",
    "random.shuffle(filenames)\n",
    "split = int(len(filenames) * VALIDATION_SPLIT)\n",
    "training_filenames = filenames[split:]\n",
    "validation_filenames = filenames[:split]\n",
    "print(\"Pattern matches {} data files. Splitting dataset into {} training files and {} validation files\".format(len(filenames), len(training_filenames), len(validation_filenames)))\n",
    "validation_steps = int(3670 // len(filenames) * len(validation_filenames)) // BATCH_SIZE\n",
    "steps_per_epoch = int(3670 // len(filenames) * len(training_filenames)) // BATCH_SIZE\n",
    "print(\"With a batch size of {}, there will be {} batches per training epoch and {} batch(es) per validation run.\".format(BATCH_SIZE, steps_per_epoch, validation_steps))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pHsCuE2DDRqX"
   },
   "source": [
    "The following cell contains some image visualization utility functions. This code isn't essential to training or deploying the model. \n",
    "\n",
    "If you're running this from Colab the cell is hidden. You can look at the code by right clicking on the cell --> \"Form\" --> \"Show form\" if you'd like to see it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "v5hLZEH3De5-"
   },
   "outputs": [],
   "source": [
    "# @title display utilities [RUN ME]\n",
    "\n",
    "def dataset_to_numpy_util(dataset, N):\n",
    "    dataset = dataset.batch(N)\n",
    "\n",
    "    if tf.executing_eagerly():\n",
    "        #  In eager mode, iterate in the Dataset directly.\n",
    "        for images, labels in dataset:\n",
    "            numpy_images = images.numpy()\n",
    "            numpy_labels = labels.numpy()\n",
    "            break\n",
    "    else:\n",
    "        # In non-eager mode, must get the TF note that\n",
    "        # yields the nextitem and run it in a tf.Session.\n",
    "        get_next_item = dataset.make_one_shot_iterator().get_next()\n",
    "        with tf.Session() as ses:\n",
    "            numpy_images, numpy_labels = ses.run(get_next_item)\n",
    "\n",
    "    return numpy_images, numpy_labels\n",
    "\n",
    "\n",
    "def title_from_label_and_target(label, correct_label):\n",
    "    label = np.argmax(label, axis=-1)  # one-hot to class number\n",
    "    correct_label = np.argmax(correct_label, axis=-1)  # one-hot to class number\n",
    "    correct = (label == correct_label)\n",
    "    return \"{} [{}{}{}]\".format(CLASSES[label], str(correct), ', shoud be ' if not correct else '',\n",
    "                                CLASSES[correct_label] if not correct else ''), correct\n",
    "\n",
    "\n",
    "def display_one_flower(image, title, subplot, red=False):\n",
    "    plt.subplot(subplot)\n",
    "    plt.axis('off')\n",
    "    plt.imshow(image)\n",
    "    plt.title(title, fontsize=16, color='red' if red else 'black')\n",
    "    return subplot + 1\n",
    "\n",
    "\n",
    "def display_9_images_from_dataset(dataset):\n",
    "    subplot = 331\n",
    "    plt.figure(figsize=(13, 13))\n",
    "    images, labels = dataset_to_numpy_util(dataset, 9)\n",
    "    for i, image in enumerate(images):\n",
    "        title = CLASSES[np.argmax(labels[i], axis=-1)]\n",
    "        subplot = display_one_flower(image, title, subplot)\n",
    "        if i >= 8:\n",
    "            break\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.subplots_adjust(wspace=0.1, hspace=0.1)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def display_9_images_with_predictions(images, predictions, labels):\n",
    "    subplot = 331\n",
    "    plt.figure(figsize=(13, 13))\n",
    "    for i, image in enumerate(images):\n",
    "        title, correct = title_from_label_and_target(predictions[i], labels[i])\n",
    "        subplot = display_one_flower(image, title, subplot, not correct)\n",
    "        if i >= 8:\n",
    "            break\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.subplots_adjust(wspace=0.1, hspace=0.1)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def display_training_curves(training, validation, title, subplot):\n",
    "    if subplot % 10 == 1:  # set up the subplots on the first call\n",
    "        plt.subplots(figsize=(10, 10), facecolor='#F0F0F0')\n",
    "        plt.tight_layout()\n",
    "    ax = plt.subplot(subplot)\n",
    "    ax.set_facecolor('#F8F8F8')\n",
    "    ax.plot(training)\n",
    "    ax.plot(validation)\n",
    "    ax.set_title('model ' + title)\n",
    "    ax.set_ylabel(title)\n",
    "    ax.set_xlabel('epoch')\n",
    "    ax.legend(['train', 'valid.'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5lguLcIRDqj-"
   },
   "source": [
    "### Read images and labels from TFRecords\n",
    "\n",
    "In this dataset the images are stored as `TFRecords`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IiQ0GryzDvE9"
   },
   "outputs": [],
   "source": [
    "def read_tfrecord(example):\n",
    "    features = {\n",
    "        \"image\": tf.io.FixedLenFeature([], tf.string),  # tf.string means bytestring\n",
    "        \"class\": tf.io.FixedLenFeature([], tf.int64),  # shape [] means scalar\n",
    "        \"one_hot_class\": tf.io.VarLenFeature(tf.float32),\n",
    "    }\n",
    "    example = tf.io.parse_single_example(example, features)\n",
    "    image = tf.image.decode_jpeg(example['image'], channels=3)\n",
    "    image = tf.cast(image, tf.float32) / 255.0  # convert image to floats in [0, 1] range\n",
    "    image = tf.reshape(image, [*IMAGE_SIZE, 3])  # explicit size will be needed for TPU\n",
    "    one_hot_class = tf.sparse.to_dense(example['one_hot_class'])\n",
    "    one_hot_class = tf.reshape(one_hot_class, [5])\n",
    "    return image, one_hot_class\n",
    "\n",
    "\n",
    "def load_dataset(filenames):\n",
    "    # Read data from TFRecords\n",
    "\n",
    "    dataset = tf.data.Dataset.from_tensor_slices(filenames)\n",
    "    dataset = dataset.interleave(tf.data.TFRecordDataset, cycle_length=16, num_parallel_calls=AUTO)  # faster\n",
    "    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTO)\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wPcefe7uEOuT"
   },
   "source": [
    "Use the visualization utility function provided earlier to preview flower images with their labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PJfoNx-_EOKI"
   },
   "outputs": [],
   "source": [
    "display_9_images_from_dataset(load_dataset(training_filenames))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0pbkgRIJFPdJ"
   },
   "source": [
    "### Create training and validation datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jqtZupcNFRSa"
   },
   "outputs": [],
   "source": [
    "def get_batched_dataset(filenames):\n",
    "    dataset = load_dataset(filenames)\n",
    "    dataset = dataset.cache()  # This dataset fits in RAM\n",
    "    dataset = dataset.repeat()\n",
    "    dataset = dataset.shuffle(2048)\n",
    "    dataset = dataset.batch(BATCH_SIZE)\n",
    "    dataset = dataset.prefetch(AUTO)  # prefetch next batch while training (autotune prefetch buffer size)\n",
    "    # For proper ordering of map/batch/repeat/prefetch, see Dataset performance guide: https://www.tensorflow.org/guide/performance/datasets\n",
    "    return dataset\n",
    "\n",
    "\n",
    "def get_training_dataset():\n",
    "    return get_batched_dataset(training_filenames)\n",
    "\n",
    "\n",
    "def get_validation_dataset():\n",
    "    return get_batched_dataset(validation_filenames)\n",
    "\n",
    "\n",
    "some_flowers, some_labels = dataset_to_numpy_util(load_dataset(validation_filenames), 8 * 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "8zr6lj66UlMn"
   },
   "source": [
    "## Build, train, and evaluate the model\n",
    "\n",
    "This section shows how to build, train, evaluate, and get local predictions from a model by using the TF.Keras [Sequential API](https://www.tensorflow.org/guide/keras/sequential_model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "Icz22E69smnD"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras import Sequential\n",
    "from tensorflow.keras.layers import Conv2D, MaxPooling2D, Dense, GlobalAveragePooling2D, BatchNormalization\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "\n",
    "model = Sequential([\n",
    "    # Stem\n",
    "    Conv2D(kernel_size=3, filters=16, padding='same', activation='relu', input_shape=[*IMAGE_SIZE, 3]),\n",
    "    BatchNormalization(),\n",
    "    Conv2D(kernel_size=3, filters=32, padding='same', activation='relu'),\n",
    "    BatchNormalization(),\n",
    "    MaxPooling2D(pool_size=2),\n",
    "\n",
    "    # Conv Group\n",
    "    Conv2D(kernel_size=3, filters=64, padding='same', activation='relu'),\n",
    "    BatchNormalization(),\n",
    "    MaxPooling2D(pool_size=2),\n",
    "    Conv2D(kernel_size=3, filters=96, padding='same', activation='relu'),\n",
    "    BatchNormalization(),\n",
    "    MaxPooling2D(pool_size=2),\n",
    "\n",
    "    # Conv Group\n",
    "    Conv2D(kernel_size=3, filters=128, padding='same', activation='relu'),\n",
    "    BatchNormalization(),\n",
    "    MaxPooling2D(pool_size=2),\n",
    "    Conv2D(kernel_size=3, filters=128, padding='same', activation='relu'),\n",
    "    BatchNormalization(),\n",
    "\n",
    "    # 1x1 Reduction\n",
    "    Conv2D(kernel_size=1, filters=32, padding='same', activation='relu'),\n",
    "    BatchNormalization(),\n",
    "\n",
    "    # Classifier\n",
    "    GlobalAveragePooling2D(),\n",
    "    Dense(5, activation='softmax')\n",
    "])\n",
    "\n",
    "model.compile(\n",
    "    optimizer=Adam(lr=0.005, decay=0.98),\n",
    "    loss='categorical_crossentropy',\n",
    "    metrics=['accuracy'])\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z3PBnvR3FCaY"
   },
   "source": [
    "### Train the model\n",
    "\n",
    "Train this on a GPU by attaching a GPU to your CAIP notebook instance. On a CPU, it'll take ~30 minutes to run training. On a GPU, it takes ~5 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vxZryg4xmdy0"
   },
   "outputs": [],
   "source": [
    "EPOCHS = 20  # Train for 60 epochs for higher accuracy, 20 should get you ~75%\n",
    "\n",
    "with tf.device(\"/GPU:0\"):\n",
    "    history = model.fit(get_training_dataset(), steps_per_epoch=steps_per_epoch, epochs=EPOCHS,\n",
    "                        validation_data=get_validation_dataset(), validation_steps=validation_steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wJ9ks5bGGBwz"
   },
   "source": [
    "### Visualize local predictions\n",
    "\n",
    "Get predictions on your local model and visualize the images with their predicted labels, using the visualization utility function provided earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yMEsR851VDZb"
   },
   "outputs": [],
   "source": [
    "# Randomize the input so that you can execute multiple times to change results\n",
    "permutation = np.random.permutation(8 * 20)\n",
    "some_flowers, some_labels = (some_flowers[permutation], some_labels[permutation])\n",
    "\n",
    "predictions = model.predict(some_flowers, batch_size=16)\n",
    "evaluations = model.evaluate(some_flowers, some_labels, batch_size=16)\n",
    "\n",
    "print(np.array(CLASSES)[np.argmax(predictions, axis=-1)].tolist())\n",
    "print('[val_loss, val_acc]', evaluations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qzCCDL1CZFx6"
   },
   "outputs": [],
   "source": [
    "display_9_images_with_predictions(some_flowers, predictions, some_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kV_NEAQwwH0e"
   },
   "source": [
    "## Export the model as a TF 2.x SavedModel\n",
    "\n",
    "When using TensorFlow 2.x, you export the model as a `SavedModel` and load it into Cloud Storage. During export, you need to define a serving function to convert data to the format your model expects. If you send encoded data to AI Platform, your serving function ensures that the data is decoded on the model server before it is passed as input to your model.\n",
    "\n",
    "### Serving function for image data\n",
    "\n",
    "Sending base 64 encoded image data to AI Platform is more space efficient. Since this deployed model expects input data as raw bytes, you need to ensure that the b64 encoded data gets converted back to raw bytes before it is passed as input to the deployed model.\n",
    "\n",
    "To resolve this, define a serving function (`serving_fn`) and attach it to the model as a preprocessing step. Add a `@tf.function` decorator so the serving function is part of the model's graph (instead of upstream on a CPU).\n",
    "\n",
    "When you send a prediction or explanation request, the request goes to the serving function (`serving_fn`), which preprocesses the b64 bytes into raw numpy bytes (`preprocess_fn`). At this point, the data can be passed to the model (`m_call`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2CRJTE7BXPJE"
   },
   "outputs": [],
   "source": [
    "export_path = 'gs://' + BUCKET_NAME + '/explanations/mymodel'\n",
    "\n",
    "\n",
    "def _preprocess(bytes_input):\n",
    "    decoded = tf.io.decode_jpeg(bytes_input, channels=3)\n",
    "    decoded = tf.image.convert_image_dtype(decoded, tf.float32)\n",
    "    resized = tf.image.resize(decoded, size=(192, 192))\n",
    "    return resized\n",
    "\n",
    "\n",
    "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
    "def preprocess_fn(bytes_inputs):\n",
    "    with tf.device(\"cpu:0\"):\n",
    "        decoded_images = tf.map_fn(_preprocess, bytes_inputs, dtype=tf.float32)\n",
    "    return {\"numpy_inputs\": decoded_images}  # User needs to make sure the key matches model's input\n",
    "\n",
    "\n",
    "m_call = tf.function(model.call).get_concrete_function([tf.TensorSpec(shape=[None, 192, 192, 3], dtype=tf.float32, name=\"numpy_inputs\")])\n",
    "\n",
    "\n",
    "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
    "def serving_fn(bytes_inputs):\n",
    "    images = preprocess_fn(bytes_inputs)\n",
    "    prob = m_call(**images)\n",
    "    return prob\n",
    "\n",
    "\n",
    "tf.saved_model.save(model, export_path, signatures={\n",
    "    'serving_default': serving_fn,\n",
    "    'xai_preprocess': preprocess_fn,  # Required for XAI\n",
    "    'xai_model': m_call  # Required for XAI\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WUhP-1OJAMCn"
   },
   "source": [
    "## Get input and output signatures\n",
    "\n",
    "Use TensorFlow's `saved_model_cli` to inspect the model's SignatureDef. You'll use this information when you deploy your model to AI Explanations in the next section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "isw99wshAMCn"
   },
   "outputs": [],
   "source": [
    "! saved_model_cli show --dir $export_path --all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_DaJseXBAMCp"
   },
   "source": [
    "You can get the signatures of your model's input and output layers by reloading the model into memory, and querying it for the signatures corresponding to each layer.\n",
    "\n",
    "You need the signatures for the following layers:\n",
    "\n",
    "- Serving function input layer\n",
    "- Model input layer\n",
    "- Model output layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lYs5Mkb8AMCp"
   },
   "outputs": [],
   "source": [
    "loaded = tf.saved_model.load(export_path)\n",
    "\n",
    "input_name = list(loaded.signatures['xai_model'].structured_input_signature[1].keys())[0]\n",
    "print(input_name)\n",
    "\n",
    "output_name = list(loaded.signatures['xai_model'].structured_outputs.keys())[0]\n",
    "print(output_name)\n",
    "\n",
    "preprocess_name = list(loaded.signatures['xai_preprocess'].structured_input_signature[1].keys())[0]\n",
    "print(preprocess_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E8g1JW03HO5v"
   },
   "source": [
    "## Generate explanation metadata\n",
    "\n",
    "In order to deploy this model to AI Explanations, you need to create an explanation_metadata.json file with information about your model inputs, outputs, and baseline. You can use the [Explainable AI SDK](https://pypi.org/project/explainable-ai-sdk/) to generate most of the fields.\n",
    "\n",
    "For image models, using `[0,1]` as your input baseline represents black and white images. This example uses np.random to generate the baseline because the training images contain a lot of black and white (i.e. daisy petals).\n",
    "\n",
    "*Note*: for the explanation request, use the model's signature for the input and output tensors. Do not use the serving function signature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ixQHycu-ahEB"
   },
   "outputs": [],
   "source": [
    "from explainable_ai_sdk.metadata.tf.v2 import SavedModelMetadataBuilder\n",
    "# We want to explain 'xai_model' signature.\n",
    "builder = SavedModelMetadataBuilder(export_path, signature_name='xai_model')\n",
    "random_baseline = np.random.rand(192, 192, 3)\n",
    "builder.set_image_metadata(\n",
    "    'numpy_inputs',\n",
    "    input_baselines=[random_baseline.tolist()])\n",
    "builder.save_metadata(export_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0fyT3PQrH1PJ"
   },
   "source": [
    "## Deploy model to AI Explanations\n",
    "\n",
    "This section shows how to use `gcloud` to deploy the model to AI Explanations, using two different explanation methods for image models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X1Oj9C2TIM7L"
   },
   "source": [
    "### Create the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lgd4HjJTIM7P"
   },
   "outputs": [],
   "source": [
    "import datetime\n",
    "MODEL = 'flowers' + TIMESTAMP\n",
    "print(MODEL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RuG90JjQIM7S"
   },
   "outputs": [],
   "source": [
    "# Create the model if it doesn't exist yet (you only need to run this once)\n",
    "! gcloud ai-platform models create $MODEL --enable-logging --region=$REGION"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7CJycA-1IM7b"
   },
   "source": [
    "### Create explainable model versions\n",
    "\n",
    "For image models, we offer two choices for explanation methods: \n",
    "* Integrated Gradients (IG)\n",
    "* XRAI \n",
    "\n",
    "You can find more info on each method in the [documentation](TODO). You can deploy a version with both so that you can compare results. **If you already know which explanation method you'd like to use, you can deploy one version and skip the code blocks for the other method.**\n",
    "\n",
    "Creating the version will take ~5-10 minutes. Note that your first deploy may take longer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QQZgzF2I25jf"
   },
   "source": [
    "#### Deploy an Integrated Gradients model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QrpL2cJ7IM7c"
   },
   "outputs": [],
   "source": [
    "# Each time you create a version the name should be unique\n",
    "IG_VERSION = 'v_ig'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "r9zZbEAeIM7h"
   },
   "outputs": [],
   "source": [
    "! gcloud beta ai-platform versions create $IG_VERSION \\\n",
    "--model $MODEL \\\n",
    "--origin $export_path \\\n",
    "--runtime-version 2.3 \\\n",
    "--framework TENSORFLOW \\\n",
    "--python-version 3.7 \\\n",
    "--machine-type n1-standard-4 \\\n",
    "--explanation-method integrated-gradients \\\n",
    "--num-integral-steps 25 \\\n",
    "--region=$REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XORWSOl3IM7i"
   },
   "outputs": [],
   "source": [
    "# Make sure the IG model deployed correctly. State should be `READY` in the following log\n",
    "! gcloud ai-platform versions describe $IG_VERSION --model $MODEL --region=$REGION"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vgvDwajk3Zus"
   },
   "source": [
    "#### Deploy an XRAI model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PT6WwCDr3c4Z"
   },
   "outputs": [],
   "source": [
    "# Each time you create a version the name should be unique\n",
    "XRAI_VERSION = 'v_xrai'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZoRixfJp3c4f"
   },
   "outputs": [],
   "source": [
    "# Create the XRAI version with gcloud\n",
    "! gcloud beta ai-platform versions create $XRAI_VERSION \\\n",
    "--model $MODEL \\\n",
    "--origin $export_path \\\n",
    "--runtime-version 2.3 \\\n",
    "--framework TENSORFLOW \\\n",
    "--python-version 3.7 \\\n",
    "--machine-type n1-standard-4 \\\n",
    "--explanation-method xrai \\\n",
    "--num-integral-steps 25 \\\n",
    "--region=$REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "khcdj5osM6ml"
   },
   "outputs": [],
   "source": [
    "# Make sure the XRAI model deployed correctly. State should be `READY` in the following log\n",
    "! gcloud ai-platform versions describe $XRAI_VERSION --model $MODEL  --region $REGION"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MWQ72uqOIelF"
   },
   "source": [
    "## Get predictions and explanations\n",
    "\n",
    "This section shows how to prepare test images to send to your deployed model, and how to send a batch prediction request to AI Explanations.\n",
    "\n",
    "### Get and prepare test images\n",
    "\n",
    "To prepare the test images:\n",
    "\n",
    "- Download a small sample of images from the flowers dataset -- just enough for a batch prediction.\n",
    "- Resize the images to match the input shape (192, 192) of the model.\n",
    "- Save the resized images back to your bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_9p4ESRVH0rk"
   },
   "outputs": [],
   "source": [
    "# Download test flowers from public bucket\n",
    "! mkdir flowers\n",
    "! gcloud storage cp gs://flowers_model/test_flowers/* ./flowers"   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LvaUEUd0ToaV"
   },
   "outputs": [],
   "source": [
    "# Resize the images to what your model is expecting (192,192)\n",
    "test_filenames = []\n",
    "\n",
    "for i in os.listdir('flowers'):\n",
    "    img_path = 'flowers/' + i\n",
    "    with PIL.Image.open(img_path) as ex_img:\n",
    "        resize_img = ex_img.resize([192, 192])\n",
    "        resize_img.save(img_path)\n",
    "        test_filenames.append(img_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-lc5V0d6AMDG"
   },
   "source": [
    "### Format your explanation request\n",
    "\n",
    "Prepare a batch of instances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AKlWo9DkAMDG"
   },
   "outputs": [],
   "source": [
    "# Prepare your images to send to your Cloud model\n",
    "instances = []\n",
    "\n",
    "for image_path in test_filenames:\n",
    "    img_bytes = tf.io.read_file(image_path)\n",
    "    b64str = b64encode(img_bytes.numpy()).decode('utf-8')\n",
    "    instances.append({preprocess_name: {'b64': b64str}})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5-B1VlyszaSz"
   },
   "source": [
    "### Send the explanations request and visualize\n",
    "\n",
    "If you deployed both an IG and an XRAI model, you can request explanations for both models and compare the results.\n",
    "\n",
    "**If you only deployed one model above, run only the cell for that explanation method.**\n",
    "\n",
    "You can use the Explainable AI SDK to send explanation requests to your deployed model and visualize the explanations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "p_Wv0UpwAMDL"
   },
   "outputs": [],
   "source": [
    "# IG EXPLANATIONS\n",
    "remote_ig_model = explainable_ai_sdk.load_model_from_ai_platform(project=PROJECT_ID, \n",
    "                                                                 model=MODEL, \n",
    "                                                                 version=IG_VERSION, \n",
    "                                                                 region=REGION)\n",
    "ig_response = remote_ig_model.explain(instances)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1fuM7GHKUxB-"
   },
   "outputs": [],
   "source": [
    "for response in ig_response:\n",
    "    response.visualize_attributions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ivs78yzb4P-f"
   },
   "outputs": [],
   "source": [
    "# XRAI EXPLANATIONS\n",
    "remote_xrai_model = explainable_ai_sdk.load_model_from_ai_platform(project=PROJECT_ID, \n",
    "                                                                   model=MODEL, \n",
    "                                                                   version=XRAI_VERSION, \n",
    "                                                                   region=REGION)\n",
    "xrai_response = remote_xrai_model.explain(instances)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U6gbL8ImZuWy"
   },
   "outputs": [],
   "source": [
    "for response in xrai_response:\n",
    "    response.visualize_attributions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "urPApeaYvJKC"
   },
   "source": [
    "## Check explanations and baselines\n",
    "\n",
    "To better make sense of your feature attributions, you can compare them with your model's baseline. For image models, the baseline_score returned by AI Explanations is the score your model would give an image input with the baseline you specified. The baseline is different for each class in the model. Every time your model predicts tulip as the top class, you'll see the same baseline score.\n",
    "\n",
    "Earlier, you used a baseline image of np.random randomly generated values. If you'd like the baseline for your model to be solid black and white images instead, pass [0,1] as the value to input_baselines in your explanation_metadata.json file above. If the baseline_score is very close to the value of example_score, the highlighted pixels may not be meaningful.\n",
    "\n",
    "Calculate the difference between baseline_score and example_score for the three test images above.\n",
    "\n",
    "Note that the score values for classification models are probabilities: the confidence your model has in its predicted class. A score of 0.90 for tulip means your model has classified the image as a tulip with 90% confidence.\n",
    "\n",
    "The code below checks baselines for the IG model. To inspect your XRAI model, swap out the ig_response and IG_VERSION variables below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gXqosrOPwD7B"
   },
   "outputs": [],
   "source": [
    "for i, response in enumerate(ig_response):\n",
    "    attr = response.get_attribution()\n",
    "    baseline_score = attr.baseline_score\n",
    "    predicted_score = attr.example_score\n",
    "    print('Baseline score: ', baseline_score)\n",
    "    print('Predicted score: ', predicted_score)\n",
    "    print('Predicted - Baseline: ', predicted_score - baseline_score, '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RLb2JmxGxaCI"
   },
   "source": [
    "### Explain the baseline image\n",
    "\n",
    "Another way to check your baseline choice is to view explanations for this model's baseline image: an image array of randomly generated values using np.random. First, convert the same np.random baseline array generated earlier to a base64 string and preview it. This encodes the random noise as if it's a PNG image. Additionally, you must convert the byte buffer to a numpy array, because this is the format the underlying model expects for input when you send the explain request."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XFyHFzMOxf2g"
   },
   "outputs": [],
   "source": [
    "# Convert your baseline from above to a base64 string\n",
    "rand_test_img = PIL.Image.fromarray((random_baseline * 255).astype('uint8'))\n",
    "buffer = io.BytesIO()\n",
    "rand_test_img.save(buffer, format=\"PNG\")\n",
    "new_image_string = b64encode(np.asarray(buffer.getvalue())).decode(\"utf-8\")\n",
    "\n",
    "# Preview it\n",
    "plt.imshow(rand_test_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XzCPzc3HyW88"
   },
   "outputs": [],
   "source": [
    "sanity_check_img = {preprocess_name: {'b64': new_image_string}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tKqckrDwyeWy"
   },
   "source": [
    "Send the explanation request for the baseline image. (To check a baseline image for XRAI, change IG_VERSION to XRAI_VERSION below.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5X2dy0l0yfZl"
   },
   "outputs": [],
   "source": [
    "# Sanity Check explanations EXPLANATIONS\n",
    "sanity_check_response = remote_ig_model.explain([sanity_check_img])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r52LKO0Xy4MT"
   },
   "source": [
    "Visualize the explanation for your random baseline image, highlighting the pixels that contributed to the prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BGWn1YRRy41h"
   },
   "outputs": [],
   "source": [
    "sanity_check_response[0].visualize_attributions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h8y9Yob0y_qF"
   },
   "source": [
    "The difference between your model's predicted score and the baseline score for this image should be close to 0. Run the following cell to confirm. If there is a difference between these two values, try increasing the number of integral steps used when you deploy your model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a1FClGu5zAKc"
   },
   "outputs": [],
   "source": [
    "attr = sanity_check_response[0].get_attribution()\n",
    "baseline_score = attr.baseline_score\n",
    "example_score = attr.example_score\n",
    "\n",
    "print(abs(baseline_score - example_score))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "x27DXeUGzb-M"
   },
   "source": [
    "## Cleaning up\n",
    "\n",
    "To clean up all GCP resources used in this project, you can [delete the GCP\n",
    "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
    "\n",
    "Alternatively, you can clean up individual resources by running the following\n",
    "commands:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "no210oWF68Uk"
   },
   "outputs": [],
   "source": [
    "# Delete model version resource\n",
    "! gcloud ai-platform versions delete $IG_VERSION --quiet --model $MODEL\n",
    "! gcloud ai-platform versions delete $XRAI_VERSION --quiet --model $MODEL\n",
    "\n",
    "# Delete model resource\n",
    "! gcloud ai-platform models delete $MODEL --quiet\n",
    "\n",
    "# Delete Cloud Storage objects that were created\n",
    "! gcloud storage rm --recursive gs://$BUCKET_NAME"   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "3F2g4OjbJ3gZ"
   },
   "source": [
    "If your Cloud Storage bucket doesn't contain any other objects and you would like to delete it, run `gcloud storage rm --recursive gs://$BUCKET_NAME`."   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true,
    "id": "K0UXLWaBJnrY"
   },
   "source": [
    "## What's next?\n",
    "\n",
    "To learn more about AI Explanations, check out the resources here.\n",
    "\n",
    "* [AI Explanations documentation](https://cloud.google.com/ai-platform/prediction/docs/ai-explanations/overview)\n",
    "* [AI Explanations whitepaper](https://storage.googleapis.com/cloud-ai-whitepapers/AI%20Explainability%20Whitepaper.pdf)\n",
    "* [Integrated gradients paper](https://arxiv.org/abs/1703.01365)\n",
    "* [XRAI paper](https://arxiv.org/abs/1906.02825)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## License"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2021 Google LLC\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "ai-explanations-image.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "environment": {
   "name": "tf2-gpu.2-3.m65",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-3:m65"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
